"""
Parsing code for the stateful line-based output format used by the
connie2 process.  Encodes a finite-state-machine implementation,
roughly based on the OO "state" design pattern.  For a similar, albeit
more generic, approach see:
http://code.activestate.com/recipes/146262/

Worth noting: each state is entered _after_ the corresponding header
line has been encountered.  Also, states that expect to see a single
line then transition may in fact successfully parse multiple valid
lines for the sake of robustness, retaining only the last.  This is
easy enough to fix (by setting a state variable marking whether they
expect to see a content line or the next state marker, private to that
implementation instance) if needed.
"""

import re
import numpy as np
import numpy.ma as ma

def is_skippable(line):
    "Ignore blank lines and anything starting with two # characters."
    return line.strip() == '' or line[0:2] == '##'

class InvalidStateInput(Exception):
    def __init__(self, state, line):
        self.state = state
        self.line  = line
    def __str__(self):
        return 'Invalid input for state; encountered "%s" while in state %s' \
            % (self.line, str(self.state))


class StateImpl(object):
    "Abstract base class for state implementations."
    def process_line(self, state, line):
        raise NotImplementedError
    def raise_invalid_state(self, line):
        raise InvalidStateInput(self.__class__, line)

class IgnoringInitialHeader(StateImpl):
    """
    Ignores comments and version info, waits for num_particles ( the
    first info we're interested in for now) Note, pretty soon we will
    probably need to read the version and branch based on that..
    """
    def process_line(self, state, line):
        if '#NUM_PARTICLES' == line:
            state.impl = NumParticles()
        elif "#ERRORS" == line:
            state.impl = NumErrors()
        else: pass

class NumParticles(StateImpl):
    """
    Parses the number of particles in total.  Expects to see each line
    having a single float (which will be converted to an int, although
    no checking on any decimal component is done), but can have
    multiple lines to parse; the results are accumulated.  Should be
    entered when a NUM_PARTICLES line is encountered, and expects to
    transition into the num-snapshots state.
    """
    def __init__(self):
        self.floatRE = re.compile(r'\d+(?:\.\d+)?')
    def process_line(self, state, line):
        np_match = self.floatRE.match(line)
        if np_match:
            state.particles += int(float(np_match.group()))
        elif '#NUM_SNAPSHOTS' == line:
            state.impl = NumSnapshots()
        else: self.raise_invalid_state(line)

class NumSnapshots(StateImpl):
    """
    Parses the number of snapshots, which is the number of time points
    for which positions are returned to the user.  Expects to see a
    line having a single float (which will be converted to an int,
    although no checking on any decimal component is done).  Should be
    entered when a NUM_SNAPSHOTS line is encountered, and expects to
    transition into the bounding-box state.
    """
    def __init__(self):
        self.floatRE = re.compile(r'\d+(?:\.\d+)?')
    def process_line(self, state, line):
        ns_match = self.floatRE.match(line)
        if ns_match:
            state.snapshots = int(float(ns_match.group()))
        elif '#BOUNDING_BOX' == line:
            state.impl = BoundingBox()
        else: self.raise_invalid_state(line)

class BoundingBox(StateImpl):
    """
    Parses the bounding box for the results.  Expects to see a line
    containing 4 floats, which are the lower-left longitude,
    lower-left latitude, upper-right longitude, and upper-right
    latitude, in that order, separated by spaces.  Should be entered
    when a BOUNDING_BOX line is encountered, and expects to transition
    into the cell-dimensions state.  Could actually parse multiple
    bounding-box lines while in this state, in which only the last
    will be retained.
    """
    def __init__(self):
        self.bbRE = re.compile(r'(?P<lonmin>-?\d+(?:\.\d+))\s+' +
                               r'(?P<latmin>-?\d+(?:\.\d+))\s+' +
                               r'(?P<lonmax>-?\d+(?:\.\d+))\s+' +
                               r'(?P<latmax>-?\d+(?:\.\d+))'     )
    def process_line(self, state, line):
        bb_match = self.bbRE.match(line)
        if '#CELL_DIMENSIONS' == line:
            state.impl = CellDimensions()
        elif bb_match:
            state.lon_min = float(bb_match.group('lonmin'))
            state.lat_min = float(bb_match.group('latmin'))
            state.lon_max = float(bb_match.group('lonmax'))
            state.lat_max = float(bb_match.group('latmax'))
        else: self.raise_invalid_state(line)

class CellDimensions(StateImpl):
    """
    Parses the cell dimensions.  Expects to see a line containing two
    floats, which are the width and height of a single box, in
    degrees.  Should be entered when a CELL_DIMENSIONS line is
    encountered, and should transition into the time-step state.
    Could actually parse multiple cell-dimensions lines while in this
    state, in which only the last will be retained.
    """
    def __init__(self):
        self.cdRE = re.compile(r'(?P<width>\d+(?:\.\d+))\s+(?P<height>\d+(?:\.\d+))')
    def process_line(self, state, line):
        cd_match = self.cdRE.match(line)
        if '#TIME_STEP' == line:
            state.impl = TimeStep()
        elif cd_match:
            state.lon_inc = float(cd_match.group('width'))
            state.lat_inc = float(cd_match.group('height'))
        else: self.raise_invalid_state(line)

class TimeStep(StateImpl):
    """
    Parses a single line, which should be a float representing the
    time step in days.  Should be entered when a TIME_STEP line is
    encountered, and expects to transition to one of STATS, TRACKS, or
    END (really should never see END, but there may be pathological
    cases).  Might actually parse multiple time-step lines, in which
    case only the last will be used.
    """
    def __init__(self):
        self.floatRE = re.compile(r'-?\d+(?:\.\d+)?')
    def process_line(self, state, line):
        if self.floatRE.match(line):
            state.time_step = float(line)
        elif '#STATS' == line:
            state.impl = Stats(state)
        elif '#FINAL_POS' == line:
            state.impl = EndPos(state)
        elif '#TRACKS' == line:
            state.impl = Tracks()
        elif '#END' == line:
            state.impl = Finished()
        else: self.raise_invalid_state(line)

class Tracks(StateImpl):
    """
    Parses a collection of lines, each describing the trace of a
    single particle.  Each line starts with an ISO8601-formatted
    time-string (in UTC time, up to minute precision), followed by a
    series of comma-separated lon-lat-depth points.  Should be entered
    when a TRACKS line is encountered, and expects to transition into
    either Stats or Finished.
    """
    def __init__(self):
        self.trackRE = re.compile(r'^(?P<time>\d{4}-\d{2}-\d{2}T\d{2}:\d{2}Z),\s+' +
                                  r'(?:-?\d+(?:\.\d+)?,\s*-?\d+(?:\.\d+)?,\s*-?\d+(?:\.\d+)?(?:,\s+|$))+')
        self.pointRE = re.compile(r'(?P<lon>-?\d+(?:\.\d+)?)' +
                                  r',\s*'                   +
                                  r'(?P<lat>-?\d+(?:\.\d+)?)' +
                                  r',\s*'                   +
                                  r'(?P<depth>-?\d+(?:\.\d+)?)' +
                                  r'(?:,\s*|$)')
    def process_line(self, state, line):
        valid_track = self.trackRE.match(line)
        if valid_track:
            # might ignore the time for now; easy to grab it later if needed
            # note in the slice that ISO8601 guarantees a fixed-width time string:
            track_points = self.pointRE.finditer(line[19:])
            track = []
            for point in track_points:
                track.append((float(point.group('lon')),
                              float(point.group('lat')),
                              float(point.group('depth'))))
            state.traces.append(track)
        elif '#STATS' == line:
            state.impl = Stats(state)
        elif '#END' == line:
            state.impl = Finished()
        else: self.raise_invalid_state(line)

class Stats(StateImpl):
    """
    Parses the box statistics.  These are a sequence of lines
    containing 3 whitespace-separated numbers; the first two are
    treated as 0-based indexes (started from the bottom left) into the
    grid formed by dividing the bounding box by the cell dimensions,
    and the second is the number of particles that pass through the
    box (thus, the statistics can then be calculated by dividing by
    the total number of particles, obtained previously).  All these
    numbers are expected to be integers, however for robustness they
    are parsed as floats before being truncated to ints.  Should be
    entered when a STATS line is encountered, and expects to
    transition into either TRACKS, FINAL_POS, or END.
    """
    def __init__(self, state):
        self.validLineRE = re.compile(r'(?P<x>\d+(?:\.\d+)?)\s+' +
                                      r'(?P<y>\d+(?:\.\d+)?)\s+' +
                                      r'(?P<num>\d+(?:\.\d+)?)'   )
    def process_line(self, state, line):
        valid = self.validLineRE.match(line)
        if valid:
            x   = int(float(valid.group('x')))
            y   = int(float(valid.group('y')))
            num = float(valid.group('num')) # keep as a float so we don't do int/int division!
            state.pushStats(x,y,num/state.particles)
        elif '#TRACKS' == line:
            state.impl = Tracks()
        elif '#FINAL_POS' == line:
            state.impl = EndPos(state)
        elif '#END' == line:
            state.impl = Finished()
        else: self.raise_invalid_state(line)

class EndPos(StateImpl):
    """
    Parses the box final position statistics.  These are a sequence of
    lines containing 3 whitespace-separated numbers; the first two are
    treated as 0-based indexes (started from the bottom left) into the
    grid formed by dividing the bounding box by the cell dimensions,
    and the second is the number of particles that finish in the box
    (thus, the statistics can then be calculated by dividing by the
    total number of particles, obtained previously).  All these
    numbers are expected to be integers, however for robustness they
    are parsed as floats before being truncated to ints.  Should be
    entered when a FINAL_POS line is encountered, and expects to
    transition into either STATS, TRACKS or END.
    """
    def __init__(self, state):
        self.validLineRE = re.compile(r'(?P<x>\d+(?:\.\d+)?)\s+' +
                                      r'(?P<y>\d+(?:\.\d+)?)\s+' +
                                      r'(?P<num>\d+(?:\.\d+)?)'   )
    def process_line(self, state, line):
        valid = self.validLineRE.match(line)
        if valid:
            x   = int(float(valid.group('x')))
            y   = int(float(valid.group('y')))
            num = float(valid.group('num')) # keep as a float so we don't do int/int division!
            state.pushFinalStats(x,y,num/state.particles)
        elif '#TRACKS' == line:
            state.impl = Tracks()
        elif '#STATS' == line:
            state.impl = Stats()
        elif '#END' == line:
            state.impl = Finished()
        else: self.raise_invalid_state(line)

class NumErrors(StateImpl):
    """Should be entered when an "#ERRORS" line is encountered.  Reads
    the number of errors, and transitions to ErrorPairs."""
    def process_line(self, state, line):
        if '#NUMBER_OF_ERRORS' == line:
            pass
        else:
            state.num_errors = int(line)
            state.impl = ErrorPairs()

class ErrorPairs(StateImpl):
    """Should be entered only from NumErrors; reads successive pairs
    of code then description lines, then transitions to Finished."""
    def __init__(self):
        self.current_code = 0   # codes start at one
        self.intre = re.compile(r'\d+')
    def process_line(self, state, line):
        if '#END' == line:
            # Should have state.num_errors = len(state.errors), but won't check that for now..
            state.impl = Finished()
        elif '#ERROR_PAIRS' == line:
            pass
        elif self.intre.match(line):
            self.current_code = int(line)
        else:
            state.errors.append((self.current_code, line))

class Finished(StateImpl):
    """
    The final state.  Should never have to parse anything, but we'll
    ignore all input from now on to be robust.  Should be entered when
    an #END line is encountered.
    """
    def process_line(self, state, line):
        pass

class State(object):
    def __init__(self):
        # the current state implementation:
        self.impl = IgnoringInitialHeader()
        # state variables to be accumulated:
        self.lon_min     = None
        self.lat_min     = None
        self.lon_max     = None
        self.lat_max     = None
        self.lon_inc     = None
        self.lat_inc     = None
        self.time_step   = None
        self.particles   = 0
        self.snapshots   = 0
        self._stats      = {} # Will be accessed as an array, via properties
        self._finalStats = {} # Will be accessed as an array, via properties
        self.traces      = [] # Currently just holds float-pairs.
        # Alternatively, errors to accumulate:
        self.num_errors  = 0
        self.errors      = [] # [(error_code, error_description)]
    def _makeMaskedArray(self, dictlist):
        # bitter experience: /occasionally/ a calc will be a infinitessimal amount below the
        # desired integral, meaning things on the greater borders can give an index out of bounds.
        # Note how we also increase the rows+cols by one; this is due to how pcolor truncates
        # the last one, so we have to accomodate it (hackishly):
        cols = round((self.lon_max - self.lon_min) / self.lon_inc) + 1
        rows = round((self.lat_max - self.lat_min) / self.lat_inc) + 1
        array = np.zeros((rows, cols))
        for (x,y), val in dictlist.iteritems():
            array[y,x] = val    # note: coords are x,y but array is row,col
        return ma.masked_equal(array, 0)
    def pushStats(self, x, y, num):
        "Should be used when parsing stats, rather than direct access (which is read-only via @properties)"
        self._stats[(x-1,y-1)] = num # convert 1-based to 0-based indexing
    def pushFinalStats(self, x, y, num):
        "Should be used when parsing end-pos stats, rather than direct access (which is read-only via @properties)"
        self._finalStats[(x-1,y-1)] = num # convert 1-based to 0-based indexing
    @property
    def stats(self):
        return self._makeMaskedArray(self._stats)
    @property
    def finalStats(self):
        return self._makeMaskedArray(self._finalStats)
    def process_line(self, line):
        self.impl.process_line(self, line)

def run_parse(filelike):
    state = State()
    for line in filelike:
        # Note: stripping spaces now (didn't used to) as well as newlines, after getting an error with input of " #END":
        line = line.strip('\n\r ')
        if is_skippable(line):
            continue
        state.process_line(line)
    return state
